from django.db import models
from django.contrib.auth.models import User
from .host import Host
import json


class SSHSession(models.Model):
    """SSH会话记录"""
    SESSION_STATUS_CHOICES = [
        ('active', '活跃'),
        ('closed', '已关闭'),
        ('timeout', '超时'),
        ('error', '错误'),
    ]
    
    session_id = models.CharField(max_length=64, unique=True, verbose_name='会话ID')
    user = models.ForeignKey(User, on_delete=models.CASCADE, verbose_name='用户')
    host = models.ForeignKey(Host, on_delete=models.CASCADE, verbose_name='主机')
    
    # 连接信息
    client_ip = models.GenericIPAddressField(verbose_name='客户端IP')
    user_agent = models.TextField(blank=True, verbose_name='用户代理')
    
    # 会话状态
    status = models.CharField(max_length=20, choices=SESSION_STATUS_CHOICES, default='active', verbose_name='状态')
    start_time = models.DateTimeField(auto_now_add=True, verbose_name='开始时间')
    end_time = models.DateTimeField(null=True, blank=True, verbose_name='结束时间')
    duration = models.IntegerField(default=0, verbose_name='持续时间(秒)')
    
    # 统计信息
    command_count = models.IntegerField(default=0, verbose_name='命令数量')
    data_size = models.BigIntegerField(default=0, verbose_name='数据大小(字节)')
    
    # 审计文件路径
    log_file_path = models.CharField(max_length=500, blank=True, verbose_name='日志文件路径')
    replay_file_path = models.CharField(max_length=500, blank=True, verbose_name='回放文件路径')
    
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    updated_at = models.DateTimeField(auto_now=True, verbose_name='更新时间')
    
    class Meta:
        db_table = 'ssh_session'
        verbose_name = 'SSH会话'
        verbose_name_plural = 'SSH会话'
        ordering = ['-start_time']
    
    def __str__(self):
        return f"{self.user.username}@{self.host.hostname} ({self.session_id})"


class SSHCommand(models.Model):
    """SSH命令记录"""
    COMMAND_TYPE_CHOICES = [
        ('input', '输入'),
        ('output', '输出'),
        ('system', '系统'),
    ]
    
    RISK_LEVEL_CHOICES = [
        ('low', '低风险'),
        ('medium', '中风险'),
        ('high', '高风险'),
        ('critical', '严重'),
    ]
    
    session = models.ForeignKey(SSHSession, on_delete=models.CASCADE, related_name='commands', verbose_name='会话')
    
    # 命令信息
    command_type = models.CharField(max_length=20, choices=COMMAND_TYPE_CHOICES, verbose_name='命令类型')
    command = models.TextField(verbose_name='命令内容')
    output = models.TextField(blank=True, verbose_name='命令输出')
    
    # 执行信息
    timestamp = models.DateTimeField(auto_now_add=True, verbose_name='执行时间')
    duration = models.FloatField(default=0, verbose_name='执行耗时(秒)')
    exit_code = models.IntegerField(null=True, blank=True, verbose_name='退出码')
    
    # 风险评估
    risk_level = models.CharField(max_length=20, choices=RISK_LEVEL_CHOICES, default='low', verbose_name='风险等级')
    is_dangerous = models.BooleanField(default=False, verbose_name='是否危险命令')
    
    # 审计信息
    working_directory = models.CharField(max_length=500, blank=True, verbose_name='工作目录')
    environment_vars = models.JSONField(default=dict, blank=True, verbose_name='环境变量')
    
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    
    class Meta:
        db_table = 'ssh_command'
        verbose_name = 'SSH命令'
        verbose_name_plural = 'SSH命令'
        ordering = ['-timestamp']
        indexes = [
            models.Index(fields=['session', 'timestamp']),
            models.Index(fields=['command_type', 'timestamp']),
            models.Index(fields=['risk_level']),
        ]
    
    def __str__(self):
        return f"{self.session.user.username}: {self.command[:50]}"


class SSHReplayData(models.Model):
    """SSH回放数据"""
    session = models.ForeignKey(SSHSession, on_delete=models.CASCADE, related_name='replay_data', verbose_name='会话')
    
    # 回放数据
    timestamp = models.FloatField(verbose_name='时间戳(相对于会话开始)')
    data_type = models.CharField(max_length=20, verbose_name='数据类型')  # input, output, resize
    data = models.TextField(verbose_name='数据内容')
    
    # 终端信息
    terminal_width = models.IntegerField(default=80, verbose_name='终端宽度')
    terminal_height = models.IntegerField(default=24, verbose_name='终端高度')
    
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    
    class Meta:
        db_table = 'ssh_replay_data'
        verbose_name = 'SSH回放数据'
        verbose_name_plural = 'SSH回放数据'
        ordering = ['timestamp']
        indexes = [
            models.Index(fields=['session', 'timestamp']),
        ]


class AuditRule(models.Model):
    """审计规则"""
    RULE_TYPE_CHOICES = [
        ('command', '命令规则'),
        ('keyword', '关键词规则'),
        ('pattern', '模式规则'),
    ]
    
    ACTION_CHOICES = [
        ('log', '记录日志'),
        ('alert', '发送告警'),
        ('block', '阻止执行'),
    ]
    
    name = models.CharField(max_length=100, verbose_name='规则名称')
    description = models.TextField(blank=True, verbose_name='规则描述')
    
    # 规则配置
    rule_type = models.CharField(max_length=20, choices=RULE_TYPE_CHOICES, verbose_name='规则类型')
    pattern = models.TextField(verbose_name='匹配模式')
    is_regex = models.BooleanField(default=False, verbose_name='是否正则表达式')
    case_sensitive = models.BooleanField(default=True, verbose_name='大小写敏感')
    
    # 动作配置
    action = models.CharField(max_length=20, choices=ACTION_CHOICES, verbose_name='触发动作')
    risk_level = models.CharField(max_length=20, choices=SSHCommand.RISK_LEVEL_CHOICES, default='medium', verbose_name='风险等级')
    
    # 状态
    is_active = models.BooleanField(default=True, verbose_name='是否启用')
    priority = models.IntegerField(default=0, verbose_name='优先级')
    
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    updated_at = models.DateTimeField(auto_now=True, verbose_name='更新时间')
    
    class Meta:
        db_table = 'audit_rule'
        verbose_name = '审计规则'
        verbose_name_plural = '审计规则'
        ordering = ['-priority', 'name']
    
    def __str__(self):
        return self.name


class AuditAlert(models.Model):
    """审计告警"""
    ALERT_STATUS_CHOICES = [
        ('pending', '待处理'),
        ('confirmed', '已确认'),
        ('ignored', '已忽略'),
        ('resolved', '已解决'),
    ]
    
    session = models.ForeignKey(SSHSession, on_delete=models.CASCADE, verbose_name='会话')
    command = models.ForeignKey(SSHCommand, on_delete=models.CASCADE, verbose_name='命令')
    rule = models.ForeignKey(AuditRule, on_delete=models.CASCADE, verbose_name='触发规则')
    
    # 告警信息
    title = models.CharField(max_length=200, verbose_name='告警标题')
    description = models.TextField(verbose_name='告警描述')
    risk_level = models.CharField(max_length=20, choices=SSHCommand.RISK_LEVEL_CHOICES, verbose_name='风险等级')
    
    # 处理状态
    status = models.CharField(max_length=20, choices=ALERT_STATUS_CHOICES, default='pending', verbose_name='处理状态')
    handler = models.ForeignKey(User, on_delete=models.SET_NULL, null=True, blank=True, related_name='handled_alerts', verbose_name='处理人')
    handle_time = models.DateTimeField(null=True, blank=True, verbose_name='处理时间')
    handle_note = models.TextField(blank=True, verbose_name='处理备注')
    
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    
    class Meta:
        db_table = 'audit_alert'
        verbose_name = '审计告警'
        verbose_name_plural = '审计告警'
        ordering = ['-created_at']
    
    def __str__(self):
        return f"{self.title} - {self.session.user.username}"
